from Crypto.PublicKey import DSA
from Crypto.Signature import DSS
from Crypto.Hash import SHA256
import base64
from FileOperate import DataProcess
import json
import numpy as np
import os

def dataListToJson(data_list):
    result = []
    for data in data_list:
        result.append(data.tolist())
    list_str = json.dumps(result)
    return list_str

def jsonToNdarrray(list_str):
    newlist = json.loads(list_str)  # <class 'list'>
    res = []
    for d in newlist:
        t = np.array(d, dtype=np.float32)
        res.append(t)
    return res

class elgamal(object):
    def produceKey(self, user_path):
        key = DSA.generate(2048)
        pri_key = key.export_key()
        pub_key = key.publickey().export_key()

        # 字节存储，w是字符串存储
        with open(user_path + "private_key.bin", "wb") as file:
            file.write(pri_key)
        with open(user_path + "public_key.pem", "wb") as file:
            file.write(pub_key)

        return key

    def sign(self, message, key, user_path):
        hasher = SHA256.new(message.encode()) # 需要字符串转字节
        signer = DSS.new(key, 'fips-186-3')
        signature = signer.sign(hasher)
        with open(user_path + 'sign.bin', 'wb') as f:
            f.write(signature)

        # 用base64转成字符串形式
        signature = str(base64.encodebytes(signature), encoding='utf-8')
        return signature

    def verify(self, message, signature, pub_key):
        hasher = SHA256.new(message.encode())
        verifier = DSS.new(pub_key, 'fips-186-3')
        # 优先逆向解密base64成bytes
        #signature = base64.decodebytes(signature.encode(encoding='utf-8'))

        try:
            verifier.verify(hasher, signature)
            print("The message is authentic.")
        except ValueError:
            print("The message is not authentic.")


if __name__ == '__main__':
    #message = "Hello"   # 用encode将str 转化为<class 'bytes'>
    model_path = '../data/ML/sourceModel/second/user0/model0.h5'
    model_data_list = DataProcess().readModel(model_path)
    message_str = dataListToJson(model_data_list)  # <class 'str'>

    user_path = "../data/socket/signature/user0/"
    if not os.path.exists(user_path):
        os.mkdir(user_path)

    key = elgamal().produceKey(user_path)
    signature = elgamal().sign(message_str, key, user_path)
    print(signature)  # 签名大概率不要

    # 直接读取公钥
    with open(user_path + 'sign.bin', 'rb') as f:
        sign_new = bytearray(f.read())
    with open(user_path + "public_key.pem", "rb") as f:
        pub_key = DSA.import_key(f.read()) # <DsaKey @0x1cf1a7fe278 y,g,p(2048),q>

    elgamal().verify(message_str, sign_new, pub_key)


    # 通过私钥生成公钥
    with open(user_path + 'private_key.bin','rb')as x:
        private_key = DSA.import_key(x.read())
    new_pub_key_temp = private_key.publickey().export_key() # b'jshdfjhajfhkaj-===='
    new_p_key = DSA.import_key(new_pub_key_temp)  # <DsaKey @0x1c99cb4f6a0 y,g,p(2048),q>

    elgamal().verify(message_str, sign_new, new_p_key)
